

import torch
from torch import nn
from torch.utils import data # 用于打乱数据，小批量。

import random
import time

class DataBatching:
    """将数据打散，分批抽取。   
    在大模型训练，测试的时候，内存有限，不可能一次性加载全部数据，因此将数据分批进行梯度下降，测试。 
    """
    def __init__(self,batch_size,data):
        """_summary_

        Args:
            batch_size (int): 每批数据的大小。
            data (列表，数组等类型，表述数据的集合): 数据集合。
        """
        self.batch_size = batch_size
        self.data = data
        self.data_count = len(data)
        self.indices = list(range(self.data_count)) # 现在indces是正序
        # 这些样本是随机读取的，没有特定的顺序  
        random.shuffle(self.indices)  # 现在indecs是乱序
        self.current_count = 0  # 当前计数，表示已经分批几个数据了。

    def next_batch(self):
        """ 获取下一个批次的数据。

        Returns:
            _type_: 返回data中随机抽取batch_size个元素的数组（或列表）。
                    如果已经抽取完了，返回None，将self.current_count置零后才可重新开始。
        """
        max_count = min(self.current_count + self.batch_size, self.data_count) # 最大的计数不超过data的数量
        if max_count == self.data_count: return None
        batch_indices = torch.tensor(self.indices[self.current_count: max_count]) # 张量数组的形式计算速度快些。
        self.current_count = max_count  
        return self.data[batch_indices]

class Data_Set(data.Dataset): 
    """ 自定义Dataset，用于给DataLoader生成数据迭代器。会被get_data_iter引用到，在项目中一般用不到。
        要实现__getitem__和__len__的函数。
    """
    def __init__(self,sample=[],lable=[]):
        self.sample = sample
        self.lable = lable

    def __len__(self):
        return len(self.sample)
    
    def __getitem__(self, index):
        return (self.sample[index], self.lable[index])

def get_data_iter(sample_list,lable_list,batch_size=256,num_workers=4):
    """ 利用from torch.utils import data 的 data.Dataloader()，生成一个迭代器。
    Args:
        sample_list (样本列表): 列表元素是一个样本的张量数据。
        lable_list (标签列表): 列表元素是标签，可以是矢量，也可以是标量。
        batch_size (int): 每个小批量包含的样本数量
        num_workers (int):  Defaults to 4：用四个线程来加载数据。当数据量很大，或者CPU 核心很多的时候，可以适当提高线程数量。
    Returns:
        _type_: 迭代器。可以当作列表使用(只能放到循环里，不能索引)，这个列表的元素是元组：([样本],[标签])，这个元组里的样本标签列表大小是batch_size。 
                用法示例：iter = get_data_iter(...)  
                for (X,Y) in iter: 这里X,Y 就是batch_size大小的列表(或者是，数组，张量，数组或张量的形式运算速度快些)
    """
    return data.DataLoader(Data_Set(sample_list,lable_list),batch_size=batch_size,shuffle=True,num_workers=num_workers)


def try_gpu(i=0): 
    """如果存在，则返回gpu(i)，否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():  
    """返回所有可用的GPU，如果没有GPU，则返回[cpu(),]"""
    devices = [torch.device(f'cuda:{i}')
             for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

def accurate_count(y_hat, y):  
    """_summary_  计算预测结果中，预测正确的数量。
    y 和 y_hat 是长度相同的数组。其中：
    y 的元素标量，表示第几类。
    y_hat 的元素是一维矢量。这个一维矢量由神经网络计算得到，矢量元素的值表示该类的概率。

    Args:
        y_hat (标量数组):  预测分类 
        y (矢量数组):   实际分类
    Returns:
        int: 预测正确的数量。
    """
    y_hat = y_hat.argmax(axis=1) # y_hat 是个二维矩阵，[行][列]，y_hat.argmax(axis=1)，就是找出每行元素最大的值，在列的维度上降维。
    cmp = y_hat.type(y.dtype) == y # 现在y,y_hat都是数组，做相等判断之后得到的cmp也是个数组。
    return (cmp.type(y.dtype).sum())

def evaluate_accuracy_cpu(net, data_iter):  
    """_summary_ 计算在指定数据集上模型的精度，在cpu上运行。

    Args:
        net (_type_): 网络模型
        data_iter (_type_): 指定数据集，数据集格式：get_data_iter 返回的迭代器。

    Returns:
        float: 计算精度
    """
    correct_count = 0
    data_count = 0
    if isinstance(net, torch.nn.Module): # 要先确定net是神经网络的类型。
        net.eval()  # 将模型设置为评估模式，结合 with torch.no_grad(): 语句关闭梯度计算。
    with torch.no_grad():  
        for X, y in data_iter: 
            y_hat = net(X)
            correct_count += accurate_count(y_hat, y)
            data_count += y.numel()  # y.numel() 获取张量 y 的元素总个数
    return correct_count / data_count

def evaluate_accuracy(net, data_iter, device=try_gpu()): #@save
    """使用GPU计算模型在数据集上的精度"""
    if isinstance(net, nn.Module):
        net.eval()  # 设置为评估模式
        if not device: # 如果 device 为空，那么就尝试从模型的参数中获取其所在的设备作为默认的运行设备。
            device = next(iter(net.parameters())).device  
            # 先将网络参数转成迭代器的类型，然后用next函数从迭代器中取出一个变量（是一个 torch.Tensor 对象）。
            # .device是torch.Tensor的一个属性，表示变量存储在哪个设备上

    # 正确预测的数量，总预测的数量
    correct_count = 0
    data_count = 0
    with torch.no_grad():
        for X, y in data_iter:
            X = X.to(device)
            y = y.to(device)
            correct_count += accurate_count(net(X), y)
            data_count += y.numel()
    return correct_count/data_count

class Timer:
    """Record multiple running times.
    这个计时器可以随开随停，每次计时的时间保存到times列表，单位是秒。
    """
    def __init__(self):
        """Defined in :numref:`subsec_linear_model`"""
        self.times = [] 
        self.start()

    def start(self):
        """Start the timer."""
        self.tik = time.time()

    def stop(self):
        """Stop the timer and record the time in a list. 返回这次的计时时间"""
        self.times.append(time.time() - self.tik)
        return self.times[-1]

    def avg(self):
        """Return the average time."""
        return sum(self.times) / len(self.times)

    def sum(self):
        """Return the sum of time."""
        return sum(self.times)

'''
知识点：

net.train()用于将模型设置为训练模式。当调用这个方法后，模型内部的一些模块（比如 Dropout 层、Batch Normalization 层等）
会根据训练模式的要求调整自身的行为，以适应训练阶段的特点。

Dropout 是一种常用的防止过拟合的技术，在训练过程中，它会按照一定的概率（例如设置为 p）随机地将一些神经元的输出置为 0，
这样可以让模型在训练时不会过于依赖某些特定的神经元，增强模型的泛化能力。

Batch Normalization 旨在对神经网络每层的输入进行归一化处理，它可以加速训练过程、提高模型的稳定性以及减少对初始化的依赖等。
其操作过程大致包括对输入数据计算均值和方差，然后进行归一化，并通过可学习的缩放和平移参数进行调整。
例如，在一个批次的图像数据经过 Batch Normalization 层时，该层会统计这批图像数据在每个通道上的均值和方差，
基于这些统计量来对数据进行归一化处理，同时根据梯度下降等优化算法不断调整缩放因子和偏移量，使得模型能够学习到合适的归一化参数。


net.eval() 是将模型设置为评估（或验证、测试）模式，与 net.train() 相对应。
在 net.eval() 模式下，Dropout 层会停止工作，也就是不再随机丢弃神经元的输出，而是直接将所有神经元的输出传递下去。
对于 Batch Normalization 层，它会使用在训练过程中累积统计得到的均值和方差（通常是移动平均的方式统计的）来对数据进行归一化处理，
而不再根据当前批次的数据重新计算均值和方差，并且也不会更新其内部的可学习参数了。

'''


def train_and_test_Regression_net(net, train_iter, test_iter, num_epochs, lr, device=try_gpu()):
    """这里训练回归网络（回归相当于拟合的意思），这里使用均方损失函数，网络输出的形状与标签的形状是相同的。
       经典用途：线性回归，光流计算等。

    Args:
        net (_type_): _description_
        train_iter (_type_): _description_
        test_iter (_type_): _description_
        num_epochs (_type_): _description_
        lr (_type_): _description_
        device (_type_, optional): _description_. Defaults to try_gpu().
    """
    pass

def train_and_test_Classify_net(net, train_iter, test_iter, num_epochs, lr, device=try_gpu()):
    """__: 用GPU训练模型(在《动手学深度学习pytorch版》第六章定义)。
           这里训练测试分类网络(交叉熵损失函数，网络输出的是每一类别的概率，而标签可以只用一个标量表示)。经典用途：图像分类。
           事先定义训练集和测试集。将训练集所有样本都做一次梯度下降，作为一轮训练，每轮数据训练完成，都会在测试集上测试一下准确性。 
           train_set和test_set都是数据集的迭代器，由前面的get_data_iter函数生成，迭代器里不知道数据量有多少，函数里会统计。

           待补充功能：
           训练完成会保存训练的网络参数，并生成损失函数曲线的图片。每10轮，保存一次网络参数，避免训练大模型的时候前功尽弃。
    Args:
        net (_type_): torch的网络模型
        train_iter (_type_): 训练集的迭代器
        test_iter (_type_): 测试集的迭代器
        num_epochs (_type_): 训练的轮次
        lr (_type_): 梯度下降的步长。
        device (_type_, optional): 计算设备，默认尝试使用GPU. Defaults to try_gpu().
    """

    def init_weights(m): 
        """定义模型参数初始化函数。目前的模型参数，只有全连接层和卷积层的参数是可以训练的。
        平均池化层，最大池化层都没有参数可以训练。循环神经网络，本质也是全连接层的排列组合。
        对于可以训练的参数，

        Args:
            m (_type_): _description_
        """
        if type(m) == nn.Linear or type(m) == nn.Conv2d:
            nn.init.xavier_uniform_(m.weight)
        """Xavier 初始化方法尝试根据输入和输出神经元的数量来合理地初始化权重，
        使得每一层的输入和输出的方差尽可能保持一致，避免梯度消失和梯度爆炸。
        xavier_uniform_ 原理：设该层输入维度为 Nin,输出维度为Nout,会从 [-sqrt(6/(Nin+Nout)),sqrt(6/(Nin+Nout))]的均匀分布中初始化权重参数。 
        xavier_normal ，则会以高斯分布的方式进行抽样。均值0，方差 2/(Nin+Nout)
        """

    net.apply(init_weights)
    print('training on', device)
    net.to(device) 
    optimizer = torch.optim.SGD(net.parameters(), lr=lr) # 随机梯度下降优化器。 “Stochastic Gradient Descent” ：随机梯度下降。
    loss = nn.CrossEntropyLoss()   # 分类模型用交叉熵损失函数

    timer, num_batches = Timer(), len(train_iter)
    for epoch in range(num_epochs):
        net.train()
        epoch_loss = 0    # 本轮模型的损失值
        epoch_correct = 0 # 本轮模型在训练集上的正确数量
        epoch_count = 0   # 本轮模型在训练集上的总样本量

        for i, (X, y) in enumerate(train_iter):
            timer.start()
            optimizer.zero_grad()
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            l.backward()
            optimizer.step()
            with torch.no_grad(): # 计算本次训练前的准确数量。
                epoch_loss += l
                epoch_correct += accurate_count(y_hat,y)
                epoch_count += X.shape[0]  # X.shape[0] 就是批次大小。
                #print(f'loss {epoch_loss:.3f}, epoch_correct {epoch_correct:.3f}, epoch_count {epoch_count:.3f}')
            timer.stop()

        test_acc = evaluate_accuracy(net, test_iter)
        train_acc = epoch_correct / epoch_count
        print(f'loss {epoch_loss:.3f}, train acc {train_acc:.3f}, test acc {test_acc:.3f}',
              f'cost time: {timer.sum():.3f} sec on {str(device)}')











